"""
将json文件转换为gt_png
"""
import json
import os
import numpy as np
import cv2
from labelme import utils
import matplotlib.pyplot as plt
 
def main():
    json_file = r'F:/total/json-pian'
    save_file = r'F:/total/json-pian/'
    list = os.listdir(json_file)   # 获取json文件列表
    for i in range(0, len(list)):
        path = os.path.join(json_file, list[i])  # 获取每个json文件的绝对路径
        filename = list[i][:-5]       # 提取出.json前的字符作为文件名，以便后续保存Label图片的时候使用
        extension = list[i][-4:]
        if extension == 'json':
            if os.path.isfile(path):
                data = json.load(open(path, 'rb'))
                img = utils.image.img_b64_to_arr(data['imageData'])  # 根据'imageData'字段的字符可以得到原图像
                ## img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
                # lbl为label图片（标注的地方用类别名对应的数字来标，其他为0）lbl_names为label名和数字的对应关系字典
                lbl, lbl_names = utils.shape.labelme_shapes_to_label(img.shape, data['shapes'])   # data['shapes']是json文件中记录着标注的位置及label等信息的字段
 
                captions = ['%d: %s' % (l, name) for l, name in enumerate(lbl_names)]
                lbl_viz = utils.draw.draw_label(lbl, img, captions)
 
                mask = []
                class_id = []
                class_name = [key for key in lbl_names.keys()]
                # 删掉背景标记
                #del class_name[0]
                ## img_mask_new = np.zeros(img_gray.shape).astype(np.uint8)
                img_mask_new = np.zeros(img.shape).astype(np.uint8)
                for i in range(1, len(lbl_names)):  # 跳过第一个class（默认为背景）
                    ## img_mask_new[lbl == i] = img_gray[lbl == i]
                    img_mask_new[lbl == i] = img[lbl == i]
                #获取标记区域
                retval, im_at_fixed = cv2.threshold(img_mask_new, 0, 255, cv2.THRESH_BINARY)
                cv2.imwrite("{}{}.png".format(save_file, filename), im_at_fixed)
                '''
                for i in range(1, len(lbl_names)):  # 跳过第一个class（默认为背景）
                    if re.match('package', class_name[i]) is not None:
                        mask.append((lbl == i).astype(np.uint8))  # 解析出像素值为1的对应，对应第一个对象 mask 为0、1组成的（0为背景，1为对象）
                        class_id.append(i)  # mask与class 一一对应
                mask = np.transpose(np.asarray(mask, np.uint8), [1, 2, 0])  # 转成[h,w,instance count]
                class_id = np.asarray(class_id, np.uint8)  # [instance count,]
                for i in range(0, len(class_id)):
                    retval, im_at_fixed = cv2.threshold(mask[:, :, i], 0, 255, cv2.THRESH_BINARY)
                    cv2.imwrite("{}{}_{}.png".format(save_file, filename, i), im_at_fixed)
                '''
if __name__ == '__main__':
    main()